from pathlib import Path
from typing import Dict, List, Union
import numpy as np
import json
from .base import BaseCoinDatabase
import logging

class JsonCoinDatabase(BaseCoinDatabase):
    """基于 JSON 文件的硬币数据库实现"""
    
    def __init__(self):
        self.front_features_dict: Dict[str, np.ndarray] = {}
        self.back_features_dict: Dict[str, np.ndarray] = {}
        self.image_paths: Dict[str, Dict[str, Path]] = {}
        self.annotations_dict: Dict[str, dict] = {}

    def add_coin(self, coin_id: str, front_features: np.ndarray, back_features: np.ndarray,
                front_path: Path, back_path: Path, annotations: dict = None):
        logging.info(f"Adding coin: {coin_id}")
        self.front_features_dict[coin_id] = front_features
        self.back_features_dict[coin_id] = back_features
        self.image_paths[coin_id] = {
            'front': front_path,
            'back': back_path
        }
        if annotations:
            self.annotations_dict[coin_id] = annotations

    def find_similar_coins(self, query_front_features: np.ndarray, query_back_features: np.ndarray,
                         top_k: int = 5) -> List[Dict[str, Union[str, float, dict]]]:
        similarities = []
        for coin_id in self.front_features_dict:
            front_similarity = np.dot(query_front_features, self.front_features_dict[coin_id])
            back_similarity = np.dot(query_back_features, self.back_features_dict[coin_id])
            combined_similarity = (front_similarity + back_similarity) / 2

            result = {
                'coin_id': coin_id,
                'front_similarity': float(front_similarity),
                'back_similarity': float(back_similarity),
                'combined_similarity': float(combined_similarity),
                'front_image_path': str(self.image_paths[coin_id]['front']),
                'back_image_path': str(self.image_paths[coin_id]['back'])
            }
            if coin_id in self.annotations_dict:
                result['annotations'] = self.annotations_dict[coin_id]
            similarities.append(result)

        similarities.sort(key=lambda x: x['combined_similarity'], reverse=True)
        return similarities[:top_k]

    def save_to_file(self, file_path: str):
        data = {
            'front_features': {k: v.tolist() for k, v in self.front_features_dict.items()},
            'back_features': {k: v.tolist() for k, v in self.back_features_dict.items()},
            'image_paths': {k: {side: str(path) for side, path in paths.items()}
                          for k, paths in self.image_paths.items()},
            'annotations': self.annotations_dict
        }
        with open(file_path, 'w') as f:
            json.dump(data, f)

    def load_from_file(self, file_path: str):
        with open(file_path, 'r') as f:
            data = json.load(f)
        self.front_features_dict = {k: np.array(v) for k, v in data['front_features'].items()}
        self.back_features_dict = {k: np.array(v) for k, v in data['back_features'].items()}
        self.image_paths = {k: {side: Path(path) for side, path in paths.items()}
                          for k, paths in data['image_paths'].items()}
        self.annotations_dict = data['annotations'] 